package cbbx_sm.decision_maker.search;

import java.io.File;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.InvalidClassException;
import java.io.ObjectInputStream;
import java.io.ObjectOutput;
import java.io.ObjectOutputStream;
import java.io.Serializable;
import java.util.ArrayList;
import java.util.Comparator;
import java.util.Hashtable;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.PriorityQueue;

import cbbx_sm.probabilistic_model.Cluster;
import cbbx_sm.probabilistic_model.ProbabilisticModel;
import cbbx_sm.utils.ExperimentManager;
import edu.wlu.cs.levy.CG.KDTree;
import edu.wlu.cs.levy.CG.KeyDuplicateException;
import edu.wlu.cs.levy.CG.KeySizeException;

public class States implements Iterable<double[]>, Serializable{
  /**
   * 
   */
	// Search for the NN taking the floor, ceiling or nn of probabilities.
	public enum NNProperty {floor, ceiling, nn};
	private static final long serialVersionUID = 1L;
  	private KDTree<Integer> statesKDT;
	private Hashtable<SchedulerState, Integer> reverseStates;
	private Hashtable<Cluster, Integer> clusterToIndex;
	private ArrayList<double[]> states;
	private Hashtable<SchedulerState, Double> probabilitiesNeeded;
    private PriorityQueue<SchedulerState> topProbabilities;

	private final int topK = ExperimentManager.topK;
	
	private class SchedulerStateComparator implements Comparator<SchedulerState>, Serializable {
		private Hashtable<SchedulerState, Double> probabilitiesNeeded;
		public SchedulerStateComparator(Hashtable<SchedulerState, Double> probabilitiesNeeded) {
			this.probabilitiesNeeded = probabilitiesNeeded;
		}
		@Override
		public int compare(SchedulerState arg0, SchedulerState arg1) {
			return probabilitiesNeeded.get(arg0).compareTo(probabilitiesNeeded.get((arg1)));
		}

	}
	public States(WorldStateTransition worldStateTransition, boolean loadPreComputedGrid){
		ArrayList<double[]> states;
		if (loadPreComputedGrid) {
			try {
				states = readUpStates(worldStateTransition.getClusterIndexHash()).states;
			} catch (InvalidClassException e) {
				e.printStackTrace();
				states = null;
				assert(states != null);
			}
		} else {
			states = computePossibleStates(worldStateTransition, 1, 0);	
		}
		prepareStates(states, worldStateTransition.getModel());
	}
		
	public States(ArrayList<double[]> states, ProbabilisticModel model){
		prepareStates(states, model);
	}
	
	private void prepareStates(ArrayList<double[]> states, ProbabilisticModel model){
		this.states = states;
		probabilitiesNeeded = new Hashtable<SchedulerState, Double>();
        Comparator<SchedulerState> comparator = new SchedulerStateComparator(probabilitiesNeeded);
        topProbabilities = 
            new PriorityQueue<SchedulerState>(topK, comparator);

		statesKDT = new KDTree<Integer>(states.iterator().next().length);
		reverseStates = new Hashtable<SchedulerState, Integer>();
		int i=0;
		for (double[] state:states){
			try {
				statesKDT.insert(state,i);
				reverseStates.put(new SchedulerState(state), i);
				i++;
			} catch (KeySizeException e) {
				// TODO Auto-generated catch block
				e.printStackTrace();
			} catch (KeyDuplicateException e) {
				// TODO Auto-generated catch block
				e.printStackTrace();
			}
		}
		/**
		 * Maintain a translation from scheduler state to clusters and back.
		 * */
	     this.clusterToIndex = new Hashtable<Cluster, Integer>();
	     for (int j=0; j < model.getFusedClusters().size(); j++){
	       this.clusterToIndex.put(model.getFusedClusters().get(j), j);
	     }
	}

	/** 
	 * Returns an id for a state if such exists, if not, returns null.
	 * 
	 * @param state
	 * @return
	 * @throws KeySizeException 
	 */
	public Integer get(double[] state){
		return reverseStates.get(new SchedulerState(state));
	}
	
	/**
	 * Returns the "compact" representation of a state as an array of doubles.
	 * 
	 * @param schedulerState the probability of motion at each cluster
	 * */
	public double[] get(Map<Cluster, Double> schedulerState) {
		double[] state = new double[this.states.get(0).length];
		int regionsFound = 0;
		for (Map.Entry<Cluster, Double> entry : schedulerState.entrySet()) {
			Integer regionId = clusterToIndex.get(entry.getKey());
			if (regionId != null){
				state[regionId] = entry.getValue();
				regionsFound ++;
			}
		}
		// Make sure that we are not dealing scheduling for the wrong clusters
		// by checking that all clusters were found in the scheduler state
		assert(regionsFound == this.states.get(0).length);
		return state;
	}
	
	public ArrayList<double[]> getStates(){
		return states;
	}
	
	public int size(){
		return states.size();
	}
	
	/** The list of all states that a given camera can be at.
	 * @return
	 */	
	protected static ArrayList<double[]> computePossibleStates(WorldStateTransition worldStateTransition, double zoomu, double upu) {
		ArrayList<double[]> states = new ArrayList<double[]>();
		/** Set the initial configuration of events at the different locations. */

		for (ScheduleConfigurationNode child: ScheduleConfigurationNode.getNextStates(null, null, worldStateTransition, false, 0d, zoomu,upu)){
			states.add(child.CiEvent);
		}

		System.out.println(" Num Zoom Regions = "+(worldStateTransition.getSchedules().length-1)+" Num States = "+states.size());

		return states;
	}
	
	@Override
	public Iterator<double[]> iterator() {
		return states.iterator();
	}

	/**
	 * Adds the most needed states to the grid and clears the "Needed" list.
	 */
	public void addTop(){	
		for (Entry<SchedulerState, Double> stateEntry : probabilitiesNeeded.entrySet()) {
			SchedulerState bestState = stateEntry.getKey();
			states.add(bestState.event);
			try {
				statesKDT.insert(bestState.event,states.size()-1);
				reverseStates.put(new SchedulerState(bestState.event),states.size()-1);
			} catch (KeySizeException e) {
				e.printStackTrace();
			} catch (KeyDuplicateException e) {
				e.printStackTrace();
			}
		}

		probabilitiesNeeded.clear();
		topProbabilities.clear();
	}

	/** finds the closest probability value and adds the probability needed to the hashtable of probs needed.
	 * 
	 * @param newEventsWithoutApproximation
	 * @param states
	 * @param floor
	 * @param log if true then the process logs the missing probability value
	 * @return
	 */
	public double[] approximateProbabilities(double[] newEventsWithoutApproximation, NNProperty nnProperty, boolean log) {		
		double[] nnState = getNNState(newEventsWithoutApproximation, nnProperty);
		
		if (log && states.size() < ExperimentManager.maxStateSize){
			SchedulerState neededState = new SchedulerState(newEventsWithoutApproximation);
			// If the state without approximation is not in the states list.
			if (!reverseStates.containsKey(neededState)) {
				Double currentDistance = probabilitiesNeeded.get(neededState);
				
				// If state is not already in the needed list
				if (currentDistance == null) {
					double bestDistance = distance(nnState, newEventsWithoutApproximation);
					
					// Get the current low distance from the head of the priority queue.
					double currentLowDistance = 
							(topProbabilities.size() == 0) ? 0 : probabilitiesNeeded.get(topProbabilities.peek());
					
					// If we are sure that we are about to add the new state
					if ((topProbabilities.size() < topK)
							|| (bestDistance > currentLowDistance)) {
						probabilitiesNeeded.put(neededState, bestDistance);
						topProbabilities.add(neededState);
						
						// If the queue goes over the limit.
						if (topProbabilities.size() > topK) {
							SchedulerState state = topProbabilities.remove();
							probabilitiesNeeded.remove(state);
						}	
					}
				}
			}
		}
		return nnState;
		
	}


	private double[] getNNState(
			double[] newEventsWithoutApproximation, NNProperty nnProperty) {
		double bestDistance = Double.MAX_VALUE;
		double[] closestState = null;
		try {
			List<Integer> candidates = null;
			switch (nnProperty) {
			case floor:
				candidates = statesKDT.range(getSearchBounds(newEventsWithoutApproximation, true), newEventsWithoutApproximation);
				break;
			case ceiling:
				candidates = statesKDT.range(newEventsWithoutApproximation, getSearchBounds(newEventsWithoutApproximation, false));
				break;
			case nn:
				candidates = statesKDT.range(getSearchBounds(newEventsWithoutApproximation, true), getSearchBounds(newEventsWithoutApproximation, false));
			}


			for (int stateId: candidates){
				double[] state = this.states.get(stateId);
				
				Double currentDistance = distance(state, newEventsWithoutApproximation);
				if (currentDistance != null && currentDistance < bestDistance){
					bestDistance = currentDistance;
					closestState = state;
				}
			}
		} catch (KeySizeException e) {
			e.printStackTrace();
		}		
		return closestState;
	}

	/**
	 * Computes the bounds of the event by approximating each dimension to the floor if {@code floor}
	 * and to the ceiling otherwise.
	 * 
	 * @param newEventsWithoutApproximation
	 * @param floor
	 * @return
	 */
	private double[] getSearchBounds(double[] newEventsWithoutApproximation, boolean floor) {
		double[] limit = new double[newEventsWithoutApproximation.length];
		for (int i = 0; i < limit.length; i++) {
			limit[i] = EventFactory.getInstance().approximateEventForProbability(newEventsWithoutApproximation[i], floor);
		}
		return limit;
	}

	private static Double distance(double[] state, double[] newEventsWithoutApproximation) {
		double distance = 0;
		for (int i=0; i<state.length; i++){
			distance += Math.abs(newEventsWithoutApproximation[i]-state[i]);
		}
		return distance;
	}

	public int getClusterIndex(Cluster cluster) {
		return this.clusterToIndex.get(cluster);
	}

	public Hashtable<Cluster, Integer> getClustersHashtable() {
		return this.clusterToIndex;
	}
	
	public void writeDown(){
		try { 
			String fileName = getFileName(clusterToIndex);
			// Serialize to a file 
			ObjectOutput out = 
				new ObjectOutputStream(new FileOutputStream(fileName)); 
			out.writeObject(this); 
			out.close(); 
		} 
		catch (IOException e) {
			e.printStackTrace();
		}
	}
	
	private String getFileName(Hashtable<Cluster, Integer> clusterToIndex) {
		// Compute the states file name
		String fileName = "";
		for (Cluster c : clusterToIndex.keySet()) {
			if (fileName.length() != 0) { 
				fileName += "_";
			}
			fileName += c.getUniqueId();
		}
		return ExperimentManager.statesInitialFolder + fileName + ".ser";
	}

	private States readUpStates(
			Hashtable<Cluster, Integer> clusterToIndex) throws InvalidClassException{
		String fileName = getFileName(clusterToIndex);

		States states = null;
		try { 
			// Deserialize from a file 
			File file = new File(fileName); 
			ObjectInputStream in = new ObjectInputStream(new FileInputStream(file)); 
			// Deserialize the object javax.swing.JButton 
			states = (States) in.readObject(); 
			in.close(); 
		}
		// If the class was changed we need to regenerate at a higher level.
		catch (InvalidClassException e) {
			throw e;
		} 
		catch (IOException e) {
			e.printStackTrace();
		} 
		catch (ClassNotFoundException e) {
			e.printStackTrace();
		}
		
		return states;
	}
}
